package cn.flexui.flexuiserver.domain;

import io.swagger.v3.oas.annotations.media.Schema;
import org.springframework.ai.chat.messages.AssistantMessage;
import org.springframework.ai.chat.messages.MessageType;
import org.springframework.ai.chat.metadata.ChatGenerationMetadata;
import org.springframework.ai.chat.metadata.ChatResponseMetadata;
import org.springframework.ai.chat.model.ChatResponse;
import org.springframework.ai.chat.model.Generation;

/**
 * @author wanmingyu
 * @create 2024/9/13 20:53
 */
@Schema(description = "聊天响应")
public record FlexChatResponse(

        @Schema(description = "id")
        String id,
        @Schema(description = "角色")
        String role,
        @Schema(description = "消息")
        String content,
        @Schema(description = "停止符号")
        String finishReason

) {

    private static final String finishType = "STOP";


    public static FlexChatResponse errorResponse(String content) {
        return new FlexChatResponse("1", MessageType.ASSISTANT.getValue(), content, finishType);
    }

    public static FlexChatResponse mapToFlexChatResponse(ChatResponse chatResponse, StringBuilder builder) {

        ChatResponseMetadata responseMetadata = chatResponse.getMetadata();
        String id = responseMetadata.getId();

        Generation result = chatResponse.getResult();
        String roleType = MessageType.ASSISTANT.getValue();
        if (result == null) {
            return new FlexChatResponse(id, roleType, "", finishType);
        }
        AssistantMessage output = result.getOutput();
        ChatGenerationMetadata metadata = result.getMetadata();
        String content = output.getContent();
        String finishReason = metadata.getFinishReason();
        if (!finishType.equalsIgnoreCase(finishReason)) {
            builder.append(content);
        }
        return new FlexChatResponse(id, roleType, content, finishReason);
    }

}
